import abc

class BaseRestorer:
    """
    图片修复抽象类
    """
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def repair(self):
        """
        修复图片隐写
        """
        pass

    def set_img_data(self, img_data):
        """
        设置图片数据
        """
        self.img_data = img_data
    
    def get_img_data(self):
        """
        获取图片数据
        """
        return self.img_data


class ImageRestorer(BaseRestorer):
    def __init__(self, img_path):
        self.img_path = img_path
        self.Restorers = {}

    def add_Restorer(self, img_type, Restorer):
        self.Restorers[img_type] = Restorer

    def repair(self):
        """
        修复图片隐写
        """
        # 读取图片
        self._read_img()
        self._analyze_img_type()

        # 选择合适的修复算法
        self._select_repair_algorithm()

        # 保存修复好的图片
        self._save_repaired_img()

    def _read_img(self):
        """
        读取图片数据
        """
        with open(self.img_path, 'rb') as f:
            self.img_data = f.read()

    def _analyze_img_type(self):
        """
        分析图片类型
        """
        # 常见图片类型的文件头签名
        image_signatures = {
            b'\xFF\xD8\xFF': 'JPEG',
            b'\x89PNG\r\n\x1a\n': 'PNG',
            b'GIF87a': 'GIF',
            b'GIF89a': 'GIF',
            b'BM': 'BMP',
            b'II*\x00': 'TIFF',  # 小端序
            b'MM\x00*': 'TIFF',  # 大端序
            b'\x00\x00\x01\x00': 'ICO',
            b'\x00\x00\x02\x00': 'CUR',
            b'8BPS': 'PSD',
            b'WEBP': 'WEBP',
        }
        for signature, img_type in image_signatures.items():
            if self.img_data.startswith(signature):
                self.img_type = img_type
                break

    def _select_repair_algorithm(self):
        """
        选择合适的修复算法
        """
        Restorer = self.Restorers.get(self.img_type)
        Restorer.set_img_data(self.img_data)
        Restorer.repair()
        self.repaired_img_data = Restorer.get_img_data()
        pass

    def _save_repaired_img(self):
        pass
